from fastapi import APIRouter, HTTPException, BackgroundTasks
from app.services.evolution_predict_service import EvolutionPredictService
from app.schemas.evolution_predict_schema import (
    AlgorithmRequest, AlgorithmResponse, AlgorithmMiddleResponse,TaskStatus
)
from app.tools.redis_client import redis_client
router = APIRouter()
service = EvolutionPredictService()

@router.post("/create", response_model=AlgorithmMiddleResponse)
async def create_prediction_task(request: AlgorithmRequest) -> AlgorithmMiddleResponse:
    """创建新的演化预测任务"""
    try:
        task_response = AlgorithmMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params
        )
        await redis_client.set_data(request.task_id, task_response.model_dump())
        return task_response
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"创建预测任务失败: {str(e)}")

@router.post("/execute/{task_id}")
async def execute_prediction_task(
    task_id: str, 
    background_tasks: BackgroundTasks
) -> AlgorithmMiddleResponse:
    """执行演化预测任务"""
    try:
        task_data = await redis_client.get_data(task_id)
        if not task_data:
            raise HTTPException(status_code=404, detail="预测任务不存在")
            
        background_tasks.add_task(
            service.run_evolution_predict,
            task_data.get("input_params"),
            task_id
        )
        
        return AlgorithmMiddleResponse(
            task_id=task_id,
            task_callback_url=task_data.get("task_callback_url"),
            task_status=TaskStatus.RUNNING,
            task_progress=0,
            input_params=task_data.get("input_params")
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"执行预测任务失败: {str(e)}")

@router.get("/result/{task_id}")
async def get_prediction_result(task_id: str) -> AlgorithmResponse:
    """获取预测结果"""
    try:
        result = await redis_client.get_data(f'{task_id}_result')
        if not result:
            raise HTTPException(status_code=404, detail="预测结果不存在")
        return AlgorithmResponse(**result)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取预测结果失败: {str(e)}") 